//
// itertool - a polynormial and type-safe iterator framework in C++.
//
// Copyright (c) 2021-2025 Feng Cong
//
// Distributed under the MIT Software License
// See accompanying file LICENSE.txt or copy at
// https://opensource.org/licenses/MIT

#pragma once

#include <itertool/itertool_fwd.h>

namespace itertool {

template <typename T, typename Function>
using MapItem = decltype(std::declval<Function>()(std::declval<T>()));

template <typename Function, typename... T>
using MapTupleItem = decltype(std::declval<Function>()(std::declval<T>()...));

namespace detail {

template <typename T, typename Function>
using MapIterableBase =
    IterableBase<MapItem<T, Function> /*New*/, T /*Old*/, Function /*Extra*/>;

template <typename T, typename Function>
class MapIterable : public MapIterableBase<T, Function> {
 public:
  using ItemType = MapItem<T, Function>;

  MapIterable(Iterator<T> it, Function func)
      : MapIterableBase<T, Function>(it, func) {}

  ItemType GetNow() override {
    return this->BaseGetExtra()(this->BaseGetIterator().GetNow());
  }

 private:
  static_assert(!std::is_void<ItemType>::value,
                "Map requires Function returns something other than void");
};

template <typename Function, typename TypeSeq, typename IndexSeq>
class MapTupleIterable;

template <typename Function, typename... T>
using MapTupleIterableBase =
    IterableBase<MapTupleItem<Function, T...> /*New*/, std::tuple<T...> /*Old*/,
                 Function /*Extra*/>;

template <typename Function, typename... T, size_t... Index>
class MapTupleIterable<Function, TypeSequence<T...>,
                       std::index_sequence<Index...>>
    : public MapTupleIterableBase<Function, T...> {
 public:
  using ItemType = MapTupleItem<Function, T...>;

  MapTupleIterable(Iterator<std::tuple<T...>> it, Function func)
      : MapTupleIterableBase<Function, T...>(it, func) {}

  ItemType GetNow() {
    auto item = this->BaseGetIterator().GetNow();
    return this->BaseGetExtra()(std::get<Index>(item)...);
  }

 private:
  static_assert(!std::is_void<ItemType>::value,
                "Map requires Function returns something other than void");
};

template <typename Function, typename... T>
using MakeMapTupleIterable = MapTupleIterable<Function, TypeSequence<T...>,
                                              std::index_sequence_for<T...>>;

}  // namespace detail

template <typename T, typename Function>
Iterator<MapItem<T, Function>> Map(Iterator<T> it, Function func) {
  return itertool::MakeIterator<detail::MapIterable<T, Function>>(it, func);
}

template <typename Function, typename... T>
Iterator<MapTupleItem<Function, T...>> Map(Iterator<std::tuple<T...>> it,
                                           Function func) {
  using IterType = detail::MakeMapTupleIterable<Function, T...>;
  return itertool::MakeIterator<IterType>(it, func);
}

}  // namespace itertool
